import torch
import torch.nn.functional as F
from torch_geometric.nn import GATConv


class GATModel(torch.nn.Module):
    # def __init__(self, num_node_features, hidden_size, out_size):
    #     super(GNNModel, self).__init__()
    #     self.conv1 = GCNConv(num_node_features, hidden_size)
    #     self.conv2 = GCNConv(hidden_size, out_size)
    #
    # def forward(self, x, edge_index):
    #     x = F.relu(self.conv1(x, edge_index))
    #     x = self.conv2(x, edge_index)
    #     return x

    def __init__(self, num_node_features, hidden_size, out_size, dropout=0.5):
        super(GATModel, self).__init__()
        self.conv1 = GATConv(num_node_features, hidden_size, dropout=dropout)
        self.conv2 = GATConv(hidden_size, out_size, dropout=dropout)

    def forward(self, x, edge_index):
        # 第一层注意力
        x = self.conv1(x, edge_index)
        # 使用LeakyReLU激活函数
        x = F.leaky_relu(x, negative_slope=0.2)
        x = F.dropout(x, p=0.5, training=self.training)
        # 第二层注意力
        x = self.conv2(x, edge_index)
        return x

    # global_mean_pool
    # 如果是图级别的预测，例如生成整个BOM，为一个图，则需要全局池化层
    # 如果是节点级别的预测，例如预测下一个组件，则不需要全局池化层
